import pandas as pd
import statsmodels.api as sm
import pingouin as pg
from pyprocessmacro import Process


def cal_f_value(cur_df: pd.DataFrame, x_cols: [], y_col: str):
    f_df = cur_df
    # 计算F值
    cs_name = '常量'  # 初始化常量
    f_df[cs_name] = 1
    # 提取自变量和因变量, 自变量，中介变量和控制变量合并为自变量
    f_x = f_df[x_cols]
    f_y = f_df[y_col]
    # 添加截距项
    f_x = sm.add_constant(f_x)
    # 创建多元线性回归模型
    f_result = sm.OLS(f_y, f_x).fit()
    cur_res = dict()
    cur_res['f'] = f_result.fvalue
    cur_res['p'] = f_result.f_pvalue
    return cur_res


def analysis_model59(cur_df: pd.DataFrame, x, y, m, w, covar):
    # 计算两列的均值
    mean_x = cur_df[x].mean()
    mean_m = cur_df[m].mean()
    mean_w = cur_df[w].mean()
    # 计算新列的值
    mult_col_name1 = 'INT1'
    cur_df[mult_col_name1] = (cur_df[x] - mean_x) * (cur_df[w] - mean_w)
    mult_col_name2 = 'INT2'
    cur_df[mult_col_name2] = (cur_df[m] - mean_m) * (cur_df[w] - mean_w)

    x2 = [x, w, mult_col_name1]
    x2.extend(covar)
    m2xw = pg.linear_regression(cur_df[x2], cur_df[m])
    m2xw_fp = cal_f_value(cur_df, x2, m)

    # x3需要的数据和x1一样
    x3 = [x, w, m, mult_col_name1, mult_col_name2]
    x3.extend(covar)
    y2xwm = pg.linear_regression(cur_df[x3], cur_df[y])
    y2xwm_fp = cal_f_value(cur_df, x3, y)

    cur_res = dict()
    cur_res['regression_m2xw'] = m2xw  # m，x对y的回归
    cur_res['regression_m2xw_fp'] = m2xw_fp  # m，x对y的回归
    cur_res['regression_y2xwm'] = y2xwm  # x对m的回归
    cur_res['regression_y2xwm_fp'] = y2xwm_fp  # x对m的回归
    return cur_res


def do_process(cur_df: pd.DataFrame, x, y, m, w, controls):
    p = Process(data=cur_df, model=59, x=x, y=y, m=[m], w=w, controls=controls, center=False, controls_in="all")
    # p.summary()
    print(p.direct_model)
    print(p.indirect_model)


if __name__ == '__main__':
    pd.set_option('expand_frame_repr', False)
    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    x_arr = ['JG']
    y_arr = ['YY']
    m_arr = ['XW']
    w_arr = ['JX']
    covar = ["性别", "年龄", "学历", "企业角色"]
    res = []
    for x in x_arr:
        for y in y_arr:
            for m in m_arr:
                for w in w_arr:
                    # 1. 线性回归
                    item_res = analysis_model59(df, x, y, m, w, covar)
                    res.append(item_res)
                    for cc in item_res:
                        print(cc)
                        print(item_res[cc])
                    print('========================================================================================')
                    # 2. process调节效应分析
                    do_process(df, x, y, m, w, covar)
    print('===================================================================================================')
